#! usr/bin/env python3
# -*- coding: utf-8 -*-

from gmpy2 import mpz, is_prime, random_state, mpz_urandomb, mpz_random, powmod, bit_length, mul, invert, t_mod, \
    next_prime, f_div
import random
import sys
import hashlib

rand = random_state(random.randrange(sys.maxsize))


class KGC(object):
    def __init__(self, bits):
        self.p_q_gen(bits)
        self.createpk()
        self.createsk()
        self.createP()

    def p_q_gen(self,bits):
        self.p=mpz(2) ** (bits - 1) + mpz_urandomb(rand, bits - 1)
        self.p = next_prime(self.p)
        self.q=mpz(2) ** (bits - 1) + mpz_urandomb(rand, bits - 1)
        self.q= next_prime(self.q)

    def createpk(self):
        self.pk=mpz(65537)

    def createsk(self):
        self.sk=invert(self.pk,(self.p-1)*(self.q-1))

    def createP(self):
        self.P=self.p*self.q

class Params(object):
    def __init__(self, pa, pb,bits):
        self.pa = pa
        self.pb = pb
        self.bits=bits

# user key class
class User(object):

    def __init__(self,uid,pk,pa,bits):
        self.pk = pk
        self.id = str(H1_hash(uid, pk, pa, bits))

    # signcryption function
    def signcrypt(self, Reciever, params, m,sk):
        w=mpz_urandomb(rand, params.bits - 1)

        h1=self.H1(self.id,Reciever.id,w,params.bits)
        t=self.T(w,Reciever.pk,params.pb)

        c=self.C(m,w,params.pb)
        h2=self.H2(self.id,t,m,c,params.bits)
        s=self.S(h2,sk,params.pa)
        Signcryption_text = (h1,t,c,s)
        return Signcryption_text

    def unsigncrypt(self, Sender, params, Signcryption_text,sk):
        h1=Signcryption_text[0]
        t=Signcryption_text[1]
        c=Signcryption_text[2]
        s=Signcryption_text[3]
        w=powmod(t,sk,params.pb)
        m=t_mod(c*invert(w,params.pb),params.pb)
        h2=self.H2(Sender.id,t,m,c,params.bits)
        h2_be=powmod(s,Sender.pk,params.pa)
        if(h1==self.H1(Sender.id,self.id,w,params.bits) and h2==h2_be):
            return m
        else:
            return False



    def H1(self,uid1,uid2,w,bits):
        return H1_hash(uid1,uid2,w,bits)

    def T(self,w,pkb,Pb):
        return powmod(w,pkb,Pb)

    def C(self,m,w,p):
        return  t_mod(m*w,p)

    def H2(self,uid1,T,m,C,bits):
        return H2_hash(uid1,T,m,C,bits)

    def S(self,h2,ska,pa):
        return powmod(h2,ska,pa)
# define the hashfunction that return 2048 bits hashvalue
def hash_func(x):
    temp1 = hashlib.sha512()
    temp2 = hashlib.sha512()
    temp3 = hashlib.sha512()
    temp4 = hashlib.sha512()
    temp1.update(x + b"1")
    temp2.update(x + b"2")
    temp3.update(x + b"3")
    temp4.update(x + b"4")
    ret_str = temp1.hexdigest() + temp2.hexdigest() + temp3.hexdigest() + temp4.hexdigest()
    return ret_str

# define the H1 hashfunction
def H1_hash(uid, x, y, bits):
    temp = hash_func(uid.encode('utf-8') + str(x).encode('utf-8') + str(y).encode('utf-8'))
    H1_hash_value = mpz(temp[: bits - 2], 16)
    return H1_hash_value

# define the H2 hashfunction
def H2_hash(uid, x, y, z, bits):
    temp = hash_func(uid.encode('utf-8') + str(x).encode('utf-8') + str(y).encode('utf-8') + str(z).encode('utf-8'))
    bits=bits/2-2
    H2_hash_value = mpz(temp[: int(bits)], 16)
    return H2_hash_value


# define the H3 hashfunction
def H3_hash(x, bits):
    temp = hash_func(str(x).encode('utf-8'))
    H3_hash_value = mpz(temp[: bits - 2], 16)
    return H3_hash_value

if __name__ == "__main__":
    alicedone= KGC(512)
    bobdone  =KGC(512)
    params=Params(alicedone.P,bobdone.P,512)
    #print(alicedone.P,bobdone.P)
    alice=User("Alice",alicedone.pk,alicedone.P,params.bits)
    bob=User("Bob",bobdone.pk,bobdone.P,params.bits)
    m="hello world!"
    Signcryption_text = alice.signcrypt(bob, params, m,alicedone.sk)
    print(Signcryption_text)
    m=bob.unsigncrypt(alice,params,Signcryption_text,bobdone.sk)
    print(m)